import gradio as gr
import numpy as np
import cv2
import pickle

# 加载模型
with open('clf2.pkl', 'rb') as file:
    model = pickle.load(file)

def classify_image(img):
    try:
        # 将PIL图像转换为OpenCV图像格式
        img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        
        # 将图像转换为灰度
        gray_img = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        
        # 调整图像大小以匹配模型训练时的尺寸
        # 假设你的模型是在32x32大小的图像上训练的
        resized_img = cv2.resize(gray_img, (32, 32))
        
        # 展平图像数组以创建一维特征向量
        flat_img = resized_img.flatten()
        
        # 将特征向量转换为模型期望的形状
        # 随机森林期望的形状是 (样本数量, 特征数量)
        # 由于我们只有一个样本，我们需要添加一个额外的维度
        input_img = flat_img.reshape(1, -1)
        
        # 使用模型进行预测
        prediction = model.predict(input_img)
        
        # 将预测结果映射回类别标签
        # 假设你的模型输出为0代表狗，1代表猫
        class_labels = {0: '狗', 1: '猫'}
        return class_labels[prediction[0]]
    except Exception as e:
        # 如果发生错误，返回错误信息
        return f"发生错误：{e}"

# 创建Gradio界面
iface = gr.Interface(
    fn=classify_image,  # 调用的函数
    inputs=gr.inputs.Image(shape=(None, None)),  # 允许上传任意尺寸的图片
    outputs="label",  # 输出为标签
)

# 启动应用程序
iface.launch()
